RDS for PostgreSQLがpg_tle 1.4.0に対応し認証をフックする処理を組み込めるようになりました

初めに

先月とはなりますがAmazon RDS for PostgreSQLがTrusted Language Extensions(pg_tle)のクライアント認証フックをサポートする旨のリリースがありました。

こちらを利用することで例えばクライアントの認証処理時に任意の処理を挟み込み追加認証を設けるようなことが可能です。

pg_tleではSQLに限らずJavaScript等の一部対応するプログラミング言語で拡張機能を作成し組み込むことができる拡張機能となります。
一般的な拡張機能としても各種言語を利用したストアドプロシージャを作成するようなものはありますがpg_tleはマネージドなシステム上でセキュリティ等の観点からファイルシステムにアクセスをさせたくない環境においてそれを制限して作成可能とする拡張機能となります。

RDS for PostgreSQLの拡張機能のバージョンで確認してみるとPostgreSQL 15.6まではv1.1.1が利用可能だったのが15.6-R2より一気に飛んでv1.4.0が利用可能となりました。

今回は元のリリースのタイトルであるクライアント認証フックを試しますがバージョンが一気に上がったことで他にも追加された機能があるので機会があればそちらも試してみようと思います。

今回やること

今回はaws_lambda拡張機能を併用しログイン成功時Lambda関数経由でにAmazon SNSによるログイン通知ができるようにしてみます。

一応開発中の17系ではログインイベントが実装されるようで将来的にはpg_tleを利用せずとも似たようなことはできそうですが先取りしてやってみようという感じです。

構成

以下のような形となります。

aws_lambdaを利用する関係でlambda用のVPCエンドポイントを用意しています。忘れがちなので注意しましょう。

今回はRDSおよびVPCエンドポイントはすでに存在する想定でそれに対して追加を行う想定での構築となります。
周辺追加リソースおよびLambda関数はSAMで実装してますので以下のリポジトリにコードを格納しておきます。

拡張機能の作成はpg_tleのリポジトリ内のドキュメントを参考にしました。

設定

RDSへのロール設定

RDSからlambda関数を実行するための権限が必要なので上記のSAMで作成されるCallLoginNotificationRdsRoleをLambda関数実行用に割り当てます。

パラメータグループの設定

以下のパラメータを設定します。なおshared_preload_libraryおよびrds.custom_dns_resolutionはタイプがStaticなため設定後は再起動が必要です。

パラメータ	値	備考
shared_preload_library	pg_tle	既存パラメータがある場合そこへの追加
pgtle.enable_clientauth	on
rds.custom_dns_resolution	1	VPCエンドポイントを利用する場合(lambda関数実行用)

独自拡張機能の追加

以下のクエリを流し拡張機能をインストールします。

pgtle.register_feature()の第一引数にログイン時に処理させたい関数を名を指定、第二引数にフックするイベント種別を指定します。

今回はログイン認証時の処理なので第二引数にはclientauthを指定します。

CREATE EXTENSION pg_tle;
-- install with aws_common
CREATE EXTENSION aws_lambda CASCADE;

--change 'postgres' to your username
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA pgtle TO postgres;

SELECT pgtle.install_extension(
  'login_notification',
  '1.0',
  'Notification login success',
$_pgtle_$
  CREATE SCHEMA login_notification;

  CREATE FUNCTION login_notification.hook_function(port pgtle.clientauth_port_subset, status integer)
  RETURNS void AS $$
    DECLARE
      execute_time timestamp;
      context json;
    BEGIN
      execute_time := now();
      
      context := '{"userName": "'|| port.user_name ||'", "timestamp": "'||  to_char(execute_time, 'YYYY-MM-DD HH24:MI:SS') || '"}';
      -- Call lambda function, if login success
      IF status = 0 THEN
        PERFORM aws_lambda.invoke(
            '{{your lambda function arn}}',
            context::json
        );
      END IF;
    END
  $$ LANGUAGE plpgsql;

  --　上記の関数を認証時のフック処理として登録
  SELECT pgtle.register_feature('login_notification.hook_function', 'clientauth');
  REVOKE ALL ON SCHEMA login_notification FROM PUBLIC;
$_pgtle_$,
'{aws_lambda}'
);

独自拡張機能の有効化

CREATE EXTENSIONで有効化します。指定する名称はpgtle.install_extension()の第一引数にした名称です。

postgres=> CREATE EXTENSION login_notification;
CREATE EXTENSION
--追加されていることを確認
postgres=> \dx
                              List of installed extensions
        Name        | Version |   Schema   |                Description
--------------------+---------+------------+--------------------------------------------
 aws_commons        | 1.2     | public     | Common data types across AWS services
 aws_lambda         | 1.0     | public     | AWS Lambda integration
 login_notification | 1.0     | public     | Notification login success
 pg_tle             | 1.4.0   | pgtle      | Trusted Language Extensions for PostgreSQL
 plpgsql            | 1.0     | pg_catalog | PL/pgSQL procedural language
(5 rows)

なお非常に重要な点となりますが、CREATE EXTENSION`実行後は残っているコネクションは確認作業を終えるまで必ず保持してください(動作確認は別ウィンドウ・クライアントで実施する)。

psql: error: connection to server at "database-2.xxxxxx.ap-northeast-1.rds.amazonaws.com" (192.168.2.251), port 5432 failed: FATAL:  invalid input syntax for type json
connection to server at "database-2.xxxxxx.ap-northeast-1.rds.amazonaws.com" (192.168.2.251), port 5432 failed: FATAL:  no pg_hba.conf entry for host "192.168.4.25", user "postgres", database "postgres", no encryption

というのも上記で定義された処理に問題がありエラーが発生すると接続自体が失敗するようになります。これのせいで泣く泣くインスタンスを一旦消す羽目になりました。

pgtle.clientauth_users_to_skip()で認証フックを適用しないユーザを指定できるみたいなので実際に利用する場合はもしもの場合に備え適用対象外のユーザを用意しておいた方が良さそうです。

動作確認

ログインして通知が来るかを確認してみます。

$ psql -h database-2.xxxxx.ap-northeast-1.rds.amazonaws.com -U postgres
Password for user postgres:
psql (14.8, server 15.6)
WARNING: psql major version 14, server major version 15.
         Some psql features might not work.
SSL connection (protocol: TLSv1.2, cipher: ECDHE-RSA-AES256-GCM-SHA384, bits: 256, compression: off)
Type "help" for help.

postgres=>

メールボックスを確認すると通知が確認できました(メッセージの内容inではなくatでは？と思いましたがリソース削除後に気づいたため時にすでに遅し)。

If the function returns a non-empty string or raises an exception, the string or exception message is interpreted as an error. clientauth will return the error to the user and fail their connection.

なお今回は認証自体に影響を与える処理ではありませんが、clientauthフックの仕様として実行される関数の返却値が空(void)もしくは空文字列の場合は認証成功扱い、例外発生もしくは文字列返却の場合は失敗とみなされるため特定条件で失敗させたいような場合はこのあたりでコントロールする形となります。

終わりに

pg_tleの認証フック利用した認証時の追加処理を行ってみました。

RDS for PostgreSQLの仕様上Zip等でパッケージングしたファイルを持ち込んでというのが難しいためpg_tle単品ではリッチな処理はやや書きづらいところではありますが、今回のようにaws_lambda拡張を併用しLambda関数に外出しすれば広く対応はできそうなので良いアイデアがある人は試してみてはいかがでしょうか。